#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : mngs_mjn_select_seq.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/19 10:23
@Description :
    根据样本比对结果，筛选指定物种对应的序列，并从最终的 FASTA 文件中导出相关序列。
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path

import pandas as pd
from Bio import SeqIO


def _parse_args() -> Namespace:
    """
    解析命令行参数。

    Returns:
        Namespace: 参数对象，包含 species、job_id、sample_id、out_file。
    """
    parser = ArgumentParser(description="Select sequences belonging to a species from final FASTA.")
    parser.add_argument("--species", type=str, required=True, help="目标物种名，与 hit_and_genus.txt 中一致。")
    parser.add_argument("--job_id", type=str, required=True, help="任务 ID，用于定位结果目录。")
    parser.add_argument("--sample_id", type=str, required=True, help="样本 ID。")
    parser.add_argument("--out_file", type=str, required=True, help="输出 FASTA 文件路径。")
    return parser.parse_args()


def main(args: Namespace) -> None:
    """
    根据指定物种名，筛选与之匹配的 seq_id，对应 final.fasta 的序列输出到 out_file。

    步骤：
        1. 读取 hit_and_genus.txt，筛选目标物种的 seq_id。
        2. 从样本的 final.fasta 中查找对应的序列。
        3. 输出到 out_file。

    Args:
        args (Namespace): 参数对象。
    """
    # 输出文件路径准备
    out_file = Path(args.out_file).absolute()
    out_file.parent.mkdir(parents=True, exist_ok=True)

    # 构建结果路径
    result_dir = Path("/data/mNGS/runmngs/result") / args.job_id / args.sample_id

    # 输入文件路径
    hit_genus_file = result_dir / f"{args.sample_id}_hit_and_genus.txt"
    final_fasta_file = result_dir / f"{args.sample_id}.final.fasta"

    # 读取比对的 seq_id - species_name 映射
    hit_genus_df = pd.read_csv(
        hit_genus_file,
        sep="\t",
        usecols=[0, 2],
        names=["seq_id", "species_name"]
    )

    # 目标 seq_id 列表
    seq_ids = set(
        hit_genus_df.query("species_name == @args.species")["seq_id"].tolist()
    )

    # 从 final.fasta 中筛选对应序列
    with open(final_fasta_file, "r") as fr, open(out_file, "w") as fw:
        filtered_records = (
            record for record in SeqIO.parse(fr, "fasta")
            if record.id in seq_ids
        )
        SeqIO.write(filtered_records, fw, "fasta")


if __name__ == "__main__":
    main(_parse_args())
